package com.es.map;

/**
 * @author 二十
 * @since 2022/1/17 10:09 下午
 */
public class TreeMap<K extends Comparable<K>, V> extends AbstractMap<K, V> {

    @Override
    public void add(K k, V v) {
        root = add(root, k, v);
    }

    private Node add(Node node, K k, V v) {
        if (node == null) {
            size++;
            return new Node(k, v);
        }
        if (k.compareTo(node.key) < 0) {
            //left
            node.left = add(node.left, k, v);
        } else if (k.compareTo(node.key) > 0) {
            //right
            node.right = add(node.right, k, v);
        }
        return node;
    }

    @Override
    public V remove(K k) {
        Node ret = get(root, k);
        if (ret.value != null) {
            root = remove(root, k);
            return ret.value;
        }
        return null;
    }

    private Node remove(Node node, K e) {
        //如果node是空，那就直接返回null
        if (node == null) {
            return null;
        }
        //如果待删除元素落在了node的左子树上，
        if (e.compareTo(node.key) < 0) {
            node.left = remove(node.left, e);
            return node;
        } else if (e.compareTo(node.key) > 0) {
            //如果待删除元素落在了node的右子树上
            node.right = remove(node.right, e);
            return node;
        } else {
            //node节点就是待删除的元素
            if (node.left == null) {
                Node right = node.right;
                node.right = null;
                size--;
                return right;
            }
            if (node.right == null) {
                Node left = node.left;
                node.left = null;
                size--;
                return left;
            }

            //否则就是node的左右节点都有数据的情况
            //这个时候需要找到比node节点大的最小节点，也就是node右子树上的最小节点
            Node min = minElement(node.right);
            //新节点的右节点就指向（移除新节点之后的右子树）
            min.right = removeMinElement(node.right);
            //新节点的左节点指向 原节点的左节点
            min.left = node.left;
            //help GC
            node.left = node.right = null;
            //返回移除指定元素后的新替换节点
            return min;
        }

    }

    private Node minElement(Node node) {
        if (node.left == null) {
            return node;
        }
        return minElement(node.left);
    }

    private Node removeMinElement(Node node) {
        //如果当前节点就是待删除节点，那就将当前节点的右子树返回，
        if (node.left == null) {
            Node p = node.right;
            node.right = null;
            size--;
            return p;
        }
        //如果当前节点不是待删除节点，继续往下递归
        node.left = removeMinElement(node.left);
        return node;
    }

    @Override
    public boolean contains(K k) {
        if (isEmpty()) {
            return false;
        }
        return contains(root, k);
    }

    private boolean contains(Node node, K k) {
        if (node == null) {
            return false;
        }
        if (node.key.equals(k)) {
            return true;
        } else if (k.compareTo(node.key) < 0) {
            return contains(node.left, k);
        } else if (k.compareTo(node.key) > 0) {
            return contains(node.right, k);
        }
        return false;
    }

    @Override
    public V get(K k) {
        return get(root, k).value;
    }

    private Node get(Node node, K k) {
        if (node == null) {
            return new Node(null, null);
        }
        if (node.key.equals(k)) {
            return node;
        } else if (k.compareTo(node.key) < 0) {
            return get(node.left, k);
        } else if (k.compareTo(node.key) > 0) {
            return get(node.right, k);
        }
        return new Node(null, null);
    }

    @Override
    public void set(K k, V v) {
        Node node = set(root, k, v);
        if (node == null) {
            add(k, v);
        }
    }

    private Node set(Node node, K k, V v) {
        if (node == null) {
            return null;
        }
        if (node.key.equals(k)) {
            node.value = v;
            return node;
        } else if (k.compareTo(node.key) < 0) {
            return set(node.left, k, v);
        } else if (k.compareTo(node.key) > 0) {
            return set(node.right, k, v);
        }
        return null;
    }

    @Override
    public int getSize() {
        return size;
    }

    @Override
    public boolean isEmpty() {
        return size == 0;
    }

    private class Node {
        private K key;

        private V value;

        private Node left, right;

        public Node(K k, V v) {
            key = k;
            value = v;
        }

        @Override
        public String toString() {
            return key + " : " + value;
        }
    }

    private int size;

    private Node root;

    @Override
    public String toString() {
        StringBuilder sb = new StringBuilder("Map:[");
        toString(root, 0, sb);
        return sb.append("]").toString();
    }

    private void toString(Node node, int depth, StringBuilder sb) {
        if (node == null) {
            return;
        }
        sb.append(node + "\t");
        toString(node.left, depth + 1, sb);
        toString(node.right, depth + 1, sb);
    }
}
